#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Datetime: 2020/11/27 10:29
# @Author  : thinkive_cfy_ide_3
# @Site    : 
# @File    : asset_ret_decomposition.py
# @Software: PyCharm 

"""
组合大类资产收益分解：分解到大类资产，不同于product_ret_decomposition是对持仓产品进行收益率分解
"""
import numpy as np
import pandas as pd

# from quant_researcher.quant.factors.factor_database.mutual_fund.fnd_nav_guess import common_functions
from quant_researcher.quant.datasource_fetch.index_api import index_price_related
from quant_researcher.quant.project_tool.common_var import RF_RATE
from quant_researcher.quant.project_tool import time_tool
from quant_researcher.quant.project_tool.time_tool import date_shifter
from quant_researcher.quant.datasource_fetch.stock_api import stock_price_related


def calc_asset_ret(asset_df_start, asset_df_end, asset_ret_df, asset_type, need_detail=False):
    """
    计算大类资产在一段时间的收益率

    :param pd.DataFrame asset_df_start: 开始时点每支证券的权重
    :param pd.DataFrame asset_df_end: 结束时点每支证券的权重
    :param pd.DataFrame asset_ret_df: 开始时点到结束时点每支证券的收益率
    :param str asset_type: 大类资产类型
    :param bool need_detail: 是否需要输出详细的中间结果
    :return:
    """
    asset_df_start = asset_df_start.rename(columns={'security_weight':'security_weight_begin'})
    asset_df_end = asset_df_end.rename(columns={'security_weight':'security_weight_end'})
    asset_df = asset_df_start.merge(asset_df_end, how='outer', on=['security_code','asset_type']).fillna(0)
    asset_df = asset_df.merge(asset_ret_df, how='left', on='security_code')
    asset_weight = (asset_df[f'security_weight_begin'] + asset_df[f'security_weight_end']) / 2
    normalized_asset_weight = asset_weight / asset_weight.sum()  # 资产内部权重需要归一化
    asset_ret = (normalized_asset_weight * asset_df['security_ret']).sum()
    if not need_detail:
        total_asset_weight = asset_weight.sum()
        result = {'asset_type':asset_type, 'asset_return':asset_ret, 'asset_weight':total_asset_weight}
        return result
    else:
        asset_df['asset_weight'] = asset_weight.sum()
        asset_df['asset_ret'] = asset_ret
        asset_df['asset_type'] = asset_type
        return asset_df


def get_active_attribution(asset_df, benchmark_asset_df):
    """
    超额收益归因

    :param pd.DataFrame asset_df: 大类资产的权重及收益
        +------------+--------------+-----------+
        | asset_type | asset_weight | asset_ret |
        +------------+--------------+-----------+
        |    cash    |    0.1283    |   0.0009  |
        |   stock    |    0.8542    |   0.1141  |
        |    bond    |    0.0009    |  -0.0001  |
        +------------+--------------+-----------+
    :param pd.DataFrame benchmark_asset_df: 基准的大类资产的权重及收益
        +------------+--------------+-------------+
        | asset_type | asset_weight | asset_ret |
        +------------+--------------+-------------+
        |    cash    |      0.0       |     0.0     |
        |   stock    |      0.7       |     0.06    |
        |    bond    |      0.3       |     0.01    |
        +------------+--------------+-------------+
    :return: pd.DataFrame
        +------------+--------------------+----------------+---------------+------------------------+--------------------+-------------------+
        | asset_type | allocation_ability | manage_ability | cross_ability | all_allocation_ability | all_manage_ability | all_other_ability |
        +------------+--------------------+----------------+---------------+------------------------+--------------------+-------------------+
        |    cash    |        0.0         |      0.0       |     0.0001    |         0.0063         |       0.0349       |      -0.0009      |
        |   stock    |       0.0093       |     0.0379     |     0.0083    |         0.0063         |       0.0349       |      -0.0009      |
        |    bond    |       -0.003       |     -0.003     |     0.003     |         0.0063         |       0.0349       |      -0.0009      |
        +------------+--------------------+----------------+---------------+------------------------+--------------------+-------------------+
    """
    fund_total_return = asset_df['total_return'][0]  # 基金总收益
    benchmark_total_return = benchmark_asset_df['total_return'][0]  # 基准的总收益
    active_ret = fund_total_return - benchmark_total_return

    df = asset_df.merge(benchmark_asset_df, how='inner', on='asset_type', suffixes=['_p','_b'])
    # 配置收益 = (资产权重 - 基准对应资产权重) * 基准对应资产收益
    df['allocation_ability'] = (df['asset_weight_p'] - df['asset_weight_b']) * df['asset_return_b']
    # 管理收益 = (资产收益 - 基准对应资产收益) * 基准对应资产权重
    df['manage_ability'] = (df['asset_return_p'] - df['asset_return_b']) * df['asset_weight_b']
    # 交叉收益 = (资产收益 - 基准对应资产收益) * (资产权重 - 基准对应资产权重)
    df['cross_ability'] = (df['asset_return_p'] - df['asset_return_b']) * (df['asset_weight_p'] - df['asset_weight_b'])

    df['all_allocation_ability'] = df['allocation_ability'].sum()
    df['all_manage_ability'] = df['manage_ability'].sum()
    df['active_ret'] = active_ret
    df['all_other_ability'] = active_ret - df['all_allocation_ability'] - df['all_manage_ability']

    return df


def get_attribution(asset_df, begin_date, end_date, benchmark, **kwargs):
    """

    :param asset_df: 大类资产的权重及收益, 必须包含cash，stock，bond三种类别的资产
        +------------+--------------+-----------+
        | asset_type | asset_weight | asset_return |
        +------------+--------------+-----------+
        |    cash    |    0.1283    |   0.0009  |
        |   stock    |    0.8542    |   0.1141  |
        |    bond    |    0.0009    |  -0.0001  |
        |    other    |    0.0009    |  -0.0001  |
        +------------+--------------+-----------+
    :param benchmark: 混合基准，目前支持 'TK70B30S', 'TK30B70S', 'TK50B50S','TK90B10S'
    :return: pd.DataFrame 组合的超额收益归因
    """
    conn = kwargs.pop('conn', None)
    assert benchmark in ['TK70B30S', 'TK30B70S', 'TK50B50S','TK90B10S'], '暂不支持非混合基准'
    total_ret = asset_df.weighted_return.sum()
    hq_start = date_shifter(begin_date,step='days',how_many=1)
    index_ret = index_price_related.get_index_return(
        index_code=[benchmark, 'CBA00301.CS', '000300'],
        start_date=hq_start,
        end_date=end_date,
        conn=conn
    )
    index_ret['end_date'] = pd.to_datetime(index_ret['end_date'])
    index_ret['end_date'] = index_ret['end_date'].dt.date
    index_ret = index_ret.rename(columns={'end_date': 'trade_date'})
    index_ret = index_ret.set_index(['trade_date', 'index_code'])['daily_return'].unstack()
    index_period_ret = (index_ret + 1).prod() - 1
    asset_df = asset_df[asset_df.asset_type != 'other']
    benchmark_asset_df = pd.DataFrame({'asset_type': ['cash', 'stock', 'bond'],
                                       'asset_weight': [0, float(benchmark[5:7])/100, float(benchmark[2:4])/100],
                                       'asset_return': [0, index_period_ret.loc['000300'], index_period_ret.loc['CBA00301.CS']]})
    # 注意，由于混合基准其实是每天都会再平衡的，因此由股债指数按比例加权得到的区间收益跟混合基准指数的区间收益并不相等，差异部分记入其他收益
    benchmark_asset_df['total_return'] = index_period_ret.loc[benchmark]
    asset_attribution = get_active_attribution(asset_df, benchmark_asset_df)
    return asset_attribution


def get_decomposition(holding_begin, holding_end, total_return, **kwargs):
    """
    
    :param pd.DataFrame holding_begin: 开始时点各大类资产的权重，格式如下
    date: 持仓日期
    asset_type：大类资产类型，目前有cash、stock、bond, other；
    security_code：证券代码；
    security_weight：证券权重

        +------------+------------+---------------+-----------------------+
        |     date      | asset_type | security_code | security_weight |
        +------------+------------+---------------+-----------------------+
        | 2019-12-31|    cash    |      cash     |         0.1121        |
        | 2019-12-31|   stock    |     000333    |         0.0968        |
        | 2019-12-31|   stock    |     000423    |         0.0192        |
        | 2019-12-31|   stock    |     000568    |         0.0879        |
        | 2019-12-31|   stock    |     000596    |         0.0816        |
        | 2019-12-31|   stock    |     000651    |         0.0971        |
        | 2019-12-31|    bond    |     050203    |         0.0162        |
        | 2019-12-31|    bond    |     150313    |         0.0049        |
        | 2019-12-31|    bond    |     150415    |         0.0049        |
        | 2019-12-31|    bond    |     170205    |         0.0049        |
        | 2019-12-31|    bond    |     190304    |         0.0049        |
        +------------+------------+---------------+-----------------------+
    :param pd.DataFrame holding_end: 结束时点各大类资产的权重，格式如下
    asset_type：大类资产类型，目前有cash、stock、bond，other；
    security_code：证券代码；
    security_weight：证券权重
        +-----------+------------+---------------+-----------------------+
        |     date    | asset_type | security_code | security_weight |
        +-----------+------------+---------------+-----------------------+
        | 2020-06-30|    cash    |      cash     |         0.1121        |
        | 2020-06-30|   stock    |     000333    |         0.0968        |
        | 2020-06-30|   stock    |     000423    |         0.0192        |
        | 2020-06-30|   stock    |     000568    |         0.0879        |
        | 2020-06-30|   stock    |     000596    |         0.0816        |
        | 2020-06-30|   stock    |     000651    |         0.0971        |
        | 2020-06-30|    bond    |     050203    |         0.0162        |
        | 2020-06-30|    bond    |     150313    |         0.0049        |
        | 2020-06-30|    bond    |     150415    |         0.0049        |
        | 2020-06-30|    bond    |     170205    |         0.0049        |
        | 2020-06-30|    bond    |     190304    |         0.0049        |
        +-----------+------------+---------------+-----------------------+
    :return: pd.Dataframe 返回组合按照持仓的绝对收益分解
    """
    conn_base = kwargs.pop('conn_base', None)
    conn_stock = kwargs.pop('conn_stock', None)

    begin = holding_begin['date'][0]
    end = holding_end['date'][0]
    hq_begin = date_shifter(begin,step='days', how_many=1)
    cash_code = 'H11025'
    cash_price = index_price_related.get_index_close(
        index_code=cash_code, start_date=hq_begin, end_date=end, conn=conn_base
    )
    if cash_price is None:
        cash_ret_df = pd.DataFrame({'security_code': [cash_code], 'security_ret': np.sqrt(RF_RATE + 1) - 1})
    else:
        cash_price = cash_price.sort_values(by='end_date')
        cash_period_ret = cash_price['close'].iloc[-1] / cash_price['close'].iloc[0] - 1
        cash_ret_df = pd.DataFrame({'security_code': [cash_code], 'security_ret': cash_period_ret})

    period_holding = pd.concat([holding_begin, holding_end])
    stock_holding = period_holding[period_holding['asset_type'] == 'stock']
    bond_holding = period_holding[period_holding['asset_type'] == 'bond']
    cash_holding = period_holding[period_holding['asset_type'] == 'cash']
    other_holding = period_holding[period_holding['asset_type'] == 'other']
    stock_list = stock_holding['security_code'].unique().tolist()
    bond_list = bond_holding['security_code'].unique().tolist()
    # 获取股票资产收益
    stock_ret_df = stock_price_related.get_stock_total_return(hq_begin, end, stock_code=stock_list, conn=conn_stock)
    stock_ret_df.columns = ['security_code','security_ret']
    # 获取债券资产收益
    bond_index_ret = index_price_related.get_index_return(
        index_code=bond_list,
        start_date=hq_begin,
        end_date=end,
        conn=conn_base
    )
    bond_index_ret['end_date'] = pd.to_datetime(bond_index_ret['end_date'])
    bond_index_ret['end_date'] = bond_index_ret['end_date'].dt.date
    bond_index_ret = bond_index_ret.rename(columns={'end_date': 'trade_date'})
    bond_index_ret = bond_index_ret.set_index(['trade_date', 'index_code'])['daily_return'].unstack()
    bond_index_ret = (bond_index_ret + 1).prod() - 1
    bond_index_ret_df = bond_index_ret.rename('security_ret').reset_index().rename(
        columns={'index_code': 'security_code'})
    cash_dict = calc_asset_ret(cash_holding[cash_holding['date']==begin], cash_holding[cash_holding['date']==end], cash_ret_df, 'cash')
    stock_dict = calc_asset_ret(stock_holding[stock_holding['date']==begin], stock_holding[stock_holding['date']==end], stock_ret_df, 'stock')
    bond_dict = calc_asset_ret(bond_holding[bond_holding['date']==begin], bond_holding[bond_holding['date']==end], bond_index_ret_df, 'bond')
    all_df = pd.DataFrame([cash_dict, stock_dict, bond_dict])
    all_df['weighted_return'] = all_df['asset_return'] * all_df['asset_weight']
    other_weighted_return = total_return - all_df.weighted_return.sum()
    other_weight_start = other_holding[other_holding['date']==begin]['security_weight'].values[0]
    other_weight_end = other_holding[other_holding['date']==end]['security_weight'].values[0]
    other_weight = (other_weight_start + other_weight_end) / 2
    other_ret = other_weighted_return / other_weight
    other_dict = {'asset_type':'other', 'asset_weight':other_weight, 'asset_return':other_ret,
                  'weighted_return': other_weighted_return}
    all_df = all_df.append(pd.Series(other_dict), ignore_index=True)
    all_df['total_return'] = total_return
    return all_df


if __name__ == '__main__':
    pass

